from turtle import forward
import torch
import torch.nn as nn

from models.custom import ResDDFNa

class MaskModule(nn.Module):

    def __init__(self, des_len, *, 
        t_min_=0.5, t_max_=1.0
    ):
        super().__init__()

        self.des_len = des_len

        self.l2_pool = nn.LPPool1d(2, self.des_len, stride=1)
        self.max_pool = nn.MaxPool1d(self.des_len, stride=1)
        self.avg_pool = nn.AvgPool1d(self.des_len, stride=1)

        self.reduce = nn.Conv2d(3, 1, kernel_size=1, bias=False)

        self.min_, self.max_ = 0.25, 0.75
        self.t_min_, self.t_max_ = t_min_, t_max_

        self.act = nn.Sigmoid()
        self.mask = nn.Hardtanh(
            min_val=self.min_, max_val=self.max_
        )

    def forward(self, x):
        b, c, h, w = x.size()
        flat = torch.swapdims(x.view(b, c, -1), 1, 2)

        l2_res = self.l2_pool(flat)
        max_res = self.max_pool(flat)
        avg_res = self.avg_pool(flat)

        res = torch.cat([l2_res, max_res, avg_res], dim=2)
        res = torch.swapdims(res, 1, 2).view(b, 3, h, w)

        res = self.reduce(res)
        confidence = self.act(res)
        mask = self.mask(confidence)

        mask = (mask-self.min_) / (self.max_-self.min_) \
            * (self.t_max_-self.t_min_) + self.t_min_

        return mask


class MaskMTanh(nn.Module):

    def __init__(self, des_len, *, 
        t_min_=0.5, t_max_=1.0
    ):
        super().__init__()

        self.des_len = des_len

        self.l2_pool = nn.LPPool1d(2, self.des_len, stride=1)
        self.max_pool = nn.MaxPool1d(self.des_len, stride=1)
        self.avg_pool = nn.AvgPool1d(self.des_len, stride=1)

        self.reduce = nn.Conv2d(3, 1, kernel_size=1, bias=False)

        self.min_, self.max_ = -1, 1
        self.t_min_, self.t_max_ = t_min_, t_max_

        self.act = nn.Tanh()

    def forward(self, x):
        b, c, h, w = x.size()
        flat = torch.swapdims(x.view(b, c, -1), 1, 2)

        l2_res = self.l2_pool(flat)
        max_res = self.max_pool(flat)
        avg_res = self.avg_pool(flat)

        res = torch.cat([l2_res, max_res, avg_res], dim=2)
        res = torch.swapdims(res, 1, 2).view(b, 3, h, w)

        res = self.reduce(res)
        mask = self.act(res)

        mask = (mask-self.min_) / (self.max_-self.min_) \
            * (self.t_max_-self.t_min_) + self.t_min_

        return mask


class MaskedResDDFNa(ResDDFNa):

    def __init__(self, out_ch=9, **kwargs):
        super().__init__(out_ch)

        self.mask = MaskModule(out_ch, **kwargs)

    def forward(self, x):

        des, norm_des = super().forward(x)
        mask = self.mask(des)

        return des, norm_des, mask


class MaskTanhResDDFNa(ResDDFNa):

    def __init__(self, out_ch=9, **kwargs):
        super().__init__(out_ch)

        self.mask = MaskMTanh(out_ch, **kwargs)

    def forward(self, x):

        des, norm_des = super().forward(x)
        mask = self.mask(des)

        return des, norm_des, mask


class MaskLearned(nn.Module):

    def __init__(self, in_ch=128):

        super().__init__()
        self.reduce = nn.Sequential(
            nn.Conv2d(in_ch, 32, kernel_size=1, bias=False), 
            nn.BatchNorm2d(32), 
            nn.ReLU()
        )
        self.mask = nn.Sequential(
            nn.Dropout(p=0.1), 
            nn.Conv2d(32, 1, kernel_size=1, bias=False), 
            nn.BatchNorm2d(1), 
            nn.Sigmoid()
        )

    def forward(self, x):

        hidden = self.reduce(x)
        mask = self.mask(hidden)

        return mask


class MaskLearnedPool(nn.Module):

    def __init__(self, in_ch=128):

        super().__init__()
        self.reduce = nn.Sequential(
            nn.Conv2d(in_ch, 32, kernel_size=1, bias=False), 
            nn.BatchNorm2d(32), 
            nn.ReLU()
        )
        self.mask = nn.Sequential(
            nn.Dropout(p=0.1), 
            nn.Conv2d(32, 8, kernel_size=1, bias=False), 
            nn.BatchNorm2d(8), 
            nn.Sigmoid()
        )
        self.pool = nn.MaxPool1d(kernel_size=8)


    def forward(self, x):
        b, c, h, w = x.size()

        hidden = self.reduce(x)
        mask = self.mask(hidden)
        mask = mask.reshape(b, 8, -1)
        mask = torch.transpose(mask, 1, 2)
        mask = self.pool(mask)
        mask = mask.reshape(b, 1, h, w)

        return mask


if __name__ == "__main__1":

    net = MaskModule(9).cuda()

    des = (torch.randint(0, 10, (2, 9, 128, 128)).float()/10).cuda()

    mask = net(des)

    print(mask)
    print(torch.max(mask))
    print(torch.min(mask))
    print(mask.shape)


if __name__ == "__main__":

    net = MaskedResDDFNa().cuda()
    img = torch.randint(0, 255, (2, 1, 128, 128)).float().cuda()

    _, _, mask = net(img)

    print(mask.shape)

